import os
import random
import shutil
import argparse
import json
import cv2
import yaml
import numpy as np
from pathlib import Path
from enum import Enum
from typing import Dict, List, Tuple, Set


class DatasetType(Enum):
    """数据集类型枚举"""
    FILTER_ONLY = "filter_only"  # 仅过滤
    YOLO_ONLY = "yolo_only"      # 生成yolo格式
    COCO_ONLY = "coco_only"      # 生成coco格式
    BOTH_FORMATS = "both_formats"  # 两种格式都生成


class DatasetProcessor:
    """
    数据集处理工具类
    支持生成YOLO格式数据集和过滤图片数据集
    """
    
    def __init__(self):
        self.image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']
        self.categories = {}
        self.image_id_counter = 1
        self.annotation_id_counter = 1
        self.category_name_to_id = {}
        self.category_id_counter = 1
    
    def filter_and_copy_images(self, image_folder, label_folder, output_folder):
        """
        将有对应标注文件的图片复制到新文件夹
        
        参数:
            image_folder: 原始图片文件夹路径
            label_folder: 标注文件文件夹路径
            output_folder: 输出文件夹路径
        """
        # 确保输出文件夹存在
        Path(output_folder).mkdir(parents=True, exist_ok=True)

        # 获取所有标注文件名（不带扩展名）
        label_files = [os.path.splitext(f)[0] for f in os.listdir(label_folder) if f.endswith('.txt')]

        # 遍历图片文件夹
        copied_count = 0
        for image_file in os.listdir(image_folder):
            # 获取不带扩展名的文件名
            image_name = os.path.splitext(image_file)[0]

            # 检查是否有对应的标注文件
            if image_name in label_files:
                # 构建完整路径
                src_path = os.path.join(image_folder, image_file)
                dst_path = os.path.join(output_folder, image_file)

                # 复制文件
                shutil.copy2(src_path, dst_path)
                copied_count += 1
                print(f"已复制: {image_file}")

        print(f"\n完成! 共复制了 {copied_count} 张图片到 {output_folder}")
        return copied_count

    def create_yolo_dataset(self, images_dir, labels_dir, output_dir='dataset',
                           train_ratio=0.8, val_ratio=0.1, test_ratio=0.1):
        """
        创建YOLO格式的数据集

        参数:
            images_dir: 图片文件夹路径
            labels_dir: 标签文件夹路径
            output_dir: 输出数据集目录
            train_ratio: 训练集比例
            val_ratio: 验证集比例
            test_ratio: 测试集比例
        """
        # 检查比例总和是否为1
        # if not (train_ratio + val_ratio + test_ratio) == 1.0:
        #     raise ValueError("训练集、验证集和测试集的比例总和必须为1.0")

        # 创建输出目录结构
        output_path = Path(output_dir)
        (output_path / 'train' / 'images').mkdir(parents=True, exist_ok=True)
        (output_path / 'train' / 'labels').mkdir(parents=True, exist_ok=True)
        (output_path / 'val' / 'images').mkdir(parents=True, exist_ok=True)
        (output_path / 'val' / 'labels').mkdir(parents=True, exist_ok=True)
        (output_path / 'test' / 'images').mkdir(parents=True, exist_ok=True)
        (output_path / 'test' / 'labels').mkdir(parents=True, exist_ok=True)

        # 获取所有图片文件(仅支持常见图片格式)
        image_files = [f for f in Path(images_dir).iterdir() if f.suffix.lower() in self.image_extensions]

        # 过滤出有对应标签文件的图片
        valid_pairs = []
        missing_labels = 0

        for img_file in image_files:
            label_file = Path(labels_dir) / (img_file.stem + '.txt')
            if label_file.exists():
                valid_pairs.append((img_file, label_file))
            else:
                print(f"警告: 图片 {img_file.name} 没有对应的标签文件")
                missing_labels += 1

        if missing_labels > 0:
            print(f"总共有 {missing_labels} 张图片没有对应的标签文件，这些图片将被忽略")

        # 随机打乱数据
        random.shuffle(valid_pairs)

        # 计算分割点
        total = len(valid_pairs)
        train_end = int(total * train_ratio)
        val_end = train_end + int(total * val_ratio)

        # 分割数据集
        train_pairs = valid_pairs[:train_end]
        val_pairs = valid_pairs[train_end:val_end]
        test_pairs = valid_pairs[val_end:]

        # 复制文件到对应目录
        def copy_files(pairs, subset):
            for img_file, label_file in pairs:
                shutil.copy(img_file, output_path / subset / 'images' / img_file.name)
                shutil.copy(label_file, output_path / subset / 'labels' / label_file.name)

        copy_files(train_pairs, 'train')
        copy_files(val_pairs, 'val')
        copy_files(test_pairs, 'test')

        # 生成dataset.yaml模板
        yaml_content = """# YOLO dataset configuration file
path: {dataset_path}  # dataset root dir
train: train/images  # train images (relative to 'path')
val: val/images      # val images (relative to 'path')
test: test/images    # test images (relative to 'path')

# Classes
names:
  0: class0
  1: class1
  2: class2
  # Add more classes as needed
"""

        with open(output_path / 'dataset.yaml', 'w') as f:
            f.write(yaml_content.format(dataset_path=str(output_path.resolve())))

        print(f"\n数据集创建完成，保存在 {output_path.resolve()}")
        print(f"数据集统计:")
        print(f"  训练集: {len(train_pairs)} 个样本")
        print(f"  验证集: {len(val_pairs)} 个样本")
        print(f"  测试集: {len(test_pairs)} 个样本")
        print(f"  总样本数: {total} (已忽略 {missing_labels} 个无标签的图片)")
        print("\n请手动修改 dataset.yaml 文件中的类别名称")
        
        return {
            'total_samples': total,
            'train_samples': len(train_pairs),
            'val_samples': len(val_pairs),
            'test_samples': len(test_pairs),
            'missing_labels': missing_labels
        }

    def create_coco_dataset(self, images_dir: str, labels_dir: str, output_dir: str, 
                          train_ratio: float = 0.8, val_ratio: float = 0.1, 
                          test_ratio: float = 0.1, class_names: List[str] = None):
        """
        创建COCO格式数据集（基于YOLO标签转换）
        
        Args:
            images_dir: 图片目录
            labels_dir: 标签目录
            output_dir: 输出目录
            train_ratio: 训练集比例
            val_ratio: 验证集比例
            test_ratio: 测试集比例
            class_names: 类别名称列表
        """
        # 验证比例
        if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-5:
            raise ValueError("训练集、验证集、测试集比例之和必须为1.0")
        
        # 创建输出目录结构
        coco_output_dir = Path(output_dir)
        coco_output_dir.mkdir(parents=True, exist_ok=True)
        
        # 创建子目录
        train_dir = coco_output_dir / "train"
        val_dir = coco_output_dir / "val"
        test_dir = coco_output_dir / "test"
        
        train_dir.mkdir(exist_ok=True)
        val_dir.mkdir(exist_ok=True)
        test_dir.mkdir(exist_ok=True)
        
        # 获取有效的图片-标签对
        valid_pairs = self._get_valid_image_label_pairs(images_dir, labels_dir)
        
        if not valid_pairs:
            print("未找到有效的图片-标签对")
            return False
        
        # 随机打乱
        random.shuffle(valid_pairs)
        
        # 分割数据集
        total_count = len(valid_pairs)
        train_count = int(total_count * train_ratio)
        val_count = int(total_count * val_ratio)
        
        train_pairs = valid_pairs[:train_count]
        val_pairs = valid_pairs[train_count:train_count + val_count]
        test_pairs = valid_pairs[train_count + val_count:]
        
        # 初始化类别映射
        if class_names:
            self._init_categories(class_names)
        else:
            # 从标签文件中提取类别
            self._extract_categories_from_labels(labels_dir)
        
        # 处理训练集
        train_coco_data = self._process_coco_split(train_pairs, train_dir, "train")
        
        # 处理验证集
        val_coco_data = self._process_coco_split(val_pairs, val_dir, "val")
        
        # 处理测试集
        test_coco_data = self._process_coco_split(test_pairs, test_dir, "test")
        
        # 保存COCO格式JSON文件
        self._save_coco_json(train_coco_data, coco_output_dir / "train.json")
        self._save_coco_json(val_coco_data, coco_output_dir / "val.json")
        self._save_coco_json(test_coco_data, coco_output_dir / "test.json")
        
        # 保存类别信息
        self._save_categories_info(coco_output_dir)
        
        print(f"COCO数据集生成完成！输出目录: {coco_output_dir}")
        print(f"训练集: {len(train_pairs)} 个样本")
        print(f"验证集: {len(val_pairs)} 个样本")
        print(f"测试集: {len(test_pairs)} 个样本")
        print(f"类别数量: {len(self.categories)}")
        
        return True

    def create_coco_dataset_from_filtered(self, images_dir: str, labels_dir: str, output_dir: str,
                                        train_ratio: float = 0.7, val_ratio: float = 0.2,
                                        test_ratio: float = 0.1, class_names: List[str] = None):
        """
        先执行过滤逻辑，然后直接生成COCO格式数据集
        
        Args:
            images_dir: 图片目录
            labels_dir: 标签目录
            output_dir: 输出目录
            train_ratio: 训练集比例
            val_ratio: 验证集比例
            test_ratio: 测试集比例
            class_names: 类别名称列表
        """
        # 验证比例
        if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-5:
            raise ValueError("训练集、验证集、测试集比例之和必须为1.0")
        
        # 创建输出目录结构
        coco_output_dir = Path(output_dir)
        coco_output_dir.mkdir(parents=True, exist_ok=True)
        
        # 创建临时过滤目录
        temp_filtered_dir = coco_output_dir / "temp_filtered"
        temp_filtered_dir.mkdir(exist_ok=True)
        
        # 第一步：执行过滤逻辑
        print("第一步：执行图片过滤...")
        filtered_pairs = self._filter_and_get_pairs(images_dir, labels_dir, temp_filtered_dir)
        
        if not filtered_pairs:
            print("过滤后未找到有效的图片-标签对")
            # 清理临时目录
            shutil.rmtree(temp_filtered_dir)
            return False
        
        print(f"过滤完成，共找到 {len(filtered_pairs)} 个有效样本")
        
        # 第二步：随机打乱
        random.shuffle(filtered_pairs)
        
        # 第三步：分割数据集
        total_count = len(filtered_pairs)
        train_count = int(total_count * train_ratio)
        val_count = int(total_count * val_ratio)
        
        train_pairs = filtered_pairs[:train_count]
        val_pairs = filtered_pairs[train_count:train_count + val_count]
        test_pairs = filtered_pairs[train_count + val_count:]
        
        # 第四步：初始化类别映射
        if class_names:
            self._init_categories(class_names)
        else:
            # 从过滤后的标签中提取类别
            self._extract_categories_from_filtered_pairs(filtered_pairs)
        
        # 第五步：创建COCO数据集目录结构
        train_dir = coco_output_dir / "train"
        val_dir = coco_output_dir / "val"
        test_dir = coco_output_dir / "test"
        
        train_dir.mkdir(exist_ok=True)
        val_dir.mkdir(exist_ok=True)
        test_dir.mkdir(exist_ok=True)
        
        # 第六步：处理各个分割集
        print("第二步：生成COCO格式数据集...")
        train_coco_data = self._process_coco_split_from_pairs(train_pairs, train_dir, "train")
        val_coco_data = self._process_coco_split_from_pairs(val_pairs, val_dir, "val")
        test_coco_data = self._process_coco_split_from_pairs(test_pairs, test_dir, "test")
        
        # 第七步：保存COCO格式JSON文件
        self._save_coco_json(train_coco_data, coco_output_dir / "train.json")
        self._save_coco_json(val_coco_data, coco_output_dir / "val.json")
        self._save_coco_json(test_coco_data, coco_output_dir / "test.json")
        
        # 第八步：保存类别信息
        self._save_categories_info(coco_output_dir)
        
        # 清理临时目录
        shutil.rmtree(temp_filtered_dir)
        
        print(f"COCO数据集生成完成！输出目录: {coco_output_dir}")
        print(f"训练集: {len(train_pairs)} 个样本")
        print(f"验证集: {len(val_pairs)} 个样本")
        print(f"测试集: {len(test_pairs)} 个样本")
        print(f"类别数量: {len(self.categories)}")
        print("注意：此方法直接基于过滤后的数据生成COCO格式，无需先生成YOLO格式")
        
        return True

    def create_both_formats(self, images_dir: str, labels_dir: str, output_dir: str,
                           train_ratio: float = 0.7, val_ratio: float = 0.2,
                           test_ratio: float = 0.1, class_names: List[str] = None):
        """
        同时生成YOLO格式和COCO格式的数据集
        
        Args:
            images_dir: 图片目录
            labels_dir: 标签目录
            output_dir: 输出目录
            train_ratio: 训练集比例
            val_ratio: 验证集比例
            test_ratio: 测试集比例
            class_names: 类别名称列表
        """
        # 验证比例
        if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-5:
            raise ValueError("训练集、验证集、测试集比例之和必须为1.0")
        
        # 创建输出目录结构
        base_output_dir = Path(output_dir)
        base_output_dir.mkdir(parents=True, exist_ok=True)
        
        # 第一步：生成YOLO格式数据集
        print("=== 第一步：生成YOLO格式数据集 ===")
        yolo_output_dir = base_output_dir / "yolo_dataset"
        yolo_result = self.create_yolo_dataset(
            images_dir=images_dir,
            labels_dir=labels_dir,
            output_dir=str(yolo_output_dir),
            train_ratio=train_ratio,
            val_ratio=val_ratio,
            test_ratio=test_ratio
        )
        
        if not yolo_result:
            print("YOLO数据集生成失败")
            return False
        
        # 第二步：生成COCO格式数据集
        print("\n=== 第二步：生成COCO格式数据集 ===")
        coco_output_dir = base_output_dir / "coco_dataset"
        coco_result = self.create_coco_dataset(
            images_dir=images_dir,
            labels_dir=labels_dir,
            output_dir=str(coco_output_dir),
            train_ratio=train_ratio,
            val_ratio=val_ratio,
            test_ratio=test_ratio,
            class_names=class_names
        )
        
        if not coco_result:
            print("COCO数据集生成失败")
            return False
        
        print(f"\n=== 两种格式数据集生成完成 ===")
        print(f"YOLO格式数据集: {yolo_output_dir}")
        print(f"COCO格式数据集: {coco_output_dir}")
        print(f"总输出目录: {base_output_dir}")
        
        return True

    def _init_categories(self, class_names: List[str]):
        """初始化类别信息"""
        self.categories = []
        self.category_name_to_id = {}
        self.category_id_counter = 1
        
        for class_name in class_names:
            category_id = self.category_id_counter
            self.categories.append({
                "id": category_id,
                "name": class_name,
                "supercategory": "object"
            })
            self.category_name_to_id[class_name] = category_id
            self.category_id_counter += 1

    def _extract_categories_from_labels(self, labels_dir: str):
        """从标签文件中提取类别信息"""
        label_files = list(Path(labels_dir).glob("*.txt"))
        unique_classes = set()
        
        for label_file in label_files:
            try:
                with open(label_file, 'r', encoding='utf-8') as f:
                    for line in f:
                        line = line.strip()
                        if line:
                            parts = line.split()
                            if len(parts) >= 5:
                                class_id = int(parts[0])
                                unique_classes.add(class_id)
            except Exception as e:
                print(f"读取标签文件 {label_file} 时出错: {e}")
        
        # 创建类别映射
        self.categories = []
        self.category_name_to_id = {}
        self.category_id_counter = 1
        
        for class_id in sorted(unique_classes):
            category_name = f"class_{class_id}"
            self.categories.append({
                "id": class_id + 1,  # COCO格式类别ID从1开始
                "name": category_name,
                "supercategory": "object"
            })
            self.category_name_to_id[category_name] = class_id + 1

    def _extract_categories_from_filtered_pairs(self, filtered_pairs: List[Tuple[Path, Path]]):
        """从过滤后的图片-标签对中提取类别信息"""
        unique_classes = set()
        
        for img_path, label_path in filtered_pairs:
            try:
                with open(label_path, 'r', encoding='utf-8') as f:
                    for line in f:
                        line = line.strip()
                        if line:
                            parts = line.split()
                            if len(parts) >= 5:
                                class_id = int(parts[0])
                                unique_classes.add(class_id)
            except Exception as e:
                print(f"读取标签文件 {label_path} 时出错: {e}")
        
        # 创建类别映射
        self.categories = []
        self.category_name_to_id = {}
        self.category_id_counter = 1
        
        for class_id in sorted(unique_classes):
            category_name = f"class_{class_id}"
            self.categories.append({
                "id": class_id + 1,  # COCO格式类别ID从1开始
                "name": category_name,
                "supercategory": "object"
            })
            self.category_name_to_id[category_name] = class_id + 1

    def _filter_and_get_pairs(self, images_dir: str, labels_dir: str, output_dir: Path) -> List[Tuple[Path, Path]]:
        """执行过滤逻辑并返回有效的图片-标签对"""
        # 获取所有图片文件
        image_files = [f for f in Path(images_dir).iterdir() if f.suffix.lower() in self.image_extensions]
        
        filtered_pairs = []
        
        for img_file in image_files:
            label_file = Path(labels_dir) / (img_file.stem + '.txt')
            if label_file.exists():
                # 复制图片到临时目录
                new_img_path = output_dir / img_file.name
                shutil.copy2(img_file, new_img_path)
                
                # 复制标签文件到临时目录
                new_label_path = output_dir / label_file.name
                shutil.copy2(label_file, new_label_path)
                
                filtered_pairs.append((img_file, label_file))
        
        return filtered_pairs

    def _process_coco_split_from_pairs(self, pairs: List[Tuple[Path, Path]], output_dir: Path, split_name: str) -> Dict:
        """处理COCO数据集的一个分割（基于图片-标签对）"""
        coco_data = {
            "info": {
                "description": "Custom Dataset",
                "year": 2023,
                "version": "1.0"
            },
            "images": [],
            "annotations": [],
            "categories": self.categories
        }
        
        for img_path, label_path in pairs:
            # 复制图片
            new_img_path = output_dir / img_path.name
            shutil.copy2(img_path, new_img_path)
            
            # 获取图片信息
            img = cv2.imread(str(img_path))
            if img is None:
                print(f"无法读取图片: {img_path}")
                continue
            
            height, width = img.shape[:2]
            
            # 创建图片信息
            image_info = {
                "id": self.image_id_counter,
                "file_name": img_path.name,
                "height": height,
                "width": width,
                "license": 1,
                "date_captured": ""
            }
            coco_data["images"].append(image_info)
            
            # 处理标签
            self._process_coco_annotations(label_path, self.image_id_counter, width, height, coco_data)
            
            self.image_id_counter += 1
        
        print(f"{split_name}分割处理完成: {len(pairs)} 个样本")
        return coco_data

    def _get_valid_image_label_pairs(self, images_dir: str, labels_dir: str) -> List[Tuple[Path, Path]]:
        """获取有效的图片-标签对"""
        valid_pairs = []
        
        # 获取所有图片文件
        image_files = [f for f in Path(images_dir).iterdir() if f.suffix.lower() in self.image_extensions]
        
        for img_file in image_files:
            label_file = Path(labels_dir) / (img_file.stem + '.txt')
            if label_file.exists():
                valid_pairs.append((img_file, label_file))
        
        return valid_pairs

    def _process_coco_split(self, pairs: List[Tuple[Path, Path]], output_dir: Path, split_name: str) -> Dict:
        """处理COCO数据集的一个分割"""
        coco_data = {
            "info": {
                "description": "Custom Dataset",
                "year": 2023,
                "version": "1.0"
            },
            "images": [],
            "annotations": [],
            "categories": self.categories
        }
        
        for img_path, label_path in pairs:
            # 复制图片
            new_img_path = output_dir / img_path.name
            shutil.copy2(img_path, new_img_path)
            
            # 获取图片信息
            img = cv2.imread(str(img_path))
            if img is None:
                print(f"无法读取图片: {img_path}")
                continue
            
            height, width = img.shape[:2]
            
            # 创建图片信息
            image_info = {
                "id": self.image_id_counter,
                "file_name": img_path.name,
                "height": height,
                "width": width,
                "license": 1,
                "date_captured": ""
            }
            coco_data["images"].append(image_info)
            
            # 处理标签
            self._process_coco_annotations(label_path, self.image_id_counter, width, height, coco_data)
            
            self.image_id_counter += 1
        
        print(f"{split_name}分割处理完成: {len(pairs)} 个样本")
        return coco_data

    def _process_coco_annotations(self, label_path: Path, image_id: int, img_width: int, img_height: int, coco_data: Dict):
        """处理COCO格式的标注"""
        try:
            with open(label_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    
                    parts = line.split()
                    if len(parts) < 5:
                        continue
                    
                    # YOLO格式: class_id x_center y_center width height
                    class_id = int(parts[0])
                    x_center = float(parts[1])
                    y_center = float(parts[2])
                    width = float(parts[3])
                    height = float(parts[4])
                    
                    # 转换为COCO格式的边界框 [x, y, width, height]
                    x = (x_center - width / 2) * img_width
                    y = (y_center - height / 2) * img_height
                    w = width * img_width
                    h = height * img_height
                    
                    # 确保边界框在图片范围内
                    x = max(0, min(x, img_width - 1))
                    y = max(0, min(y, img_height - 1))
                    w = min(w, img_width - x)
                    h = min(h, img_height - y)
                    
                    # 创建标注信息
                    annotation = {
                        "id": self.annotation_id_counter,
                        "image_id": image_id,
                        "category_id": class_id + 1,  # COCO格式类别ID从1开始
                        "bbox": [x, y, w, h],
                        "area": w * h,
                        "iscrowd": 0,
                        "segmentation": []
                    }
                    coco_data["annotations"].append(annotation)
                    self.annotation_id_counter += 1
                    
        except Exception as e:
            print(f"处理标注文件 {label_path} 时出错: {e}")

    def _save_coco_json(self, coco_data: Dict, output_path: Path):
        """保存COCO格式的JSON文件"""
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(coco_data, f, indent=2, ensure_ascii=False)
        print(f"COCO JSON文件已保存: {output_path}")

    def _save_categories_info(self, output_dir: Path):
        """保存类别信息文件"""
        categories_info = {
            "categories": self.categories,
            "category_mapping": self.category_name_to_id
        }
        
        with open(output_dir / "categories.json", 'w', encoding='utf-8') as f:
            json.dump(categories_info, f, indent=2, ensure_ascii=False)
        
        # 同时保存为YAML格式
        with open(output_dir / "categories.yaml", 'w', encoding='utf-8') as f:
            yaml.dump(categories_info, f, default_flow_style=False)

    def process_dataset(self, dataset_type: DatasetType, **kwargs):
        """
        统一的数据集处理接口
        
        Args:
            dataset_type: 数据集类型枚举
            **kwargs: 根据数据集类型传递不同的参数
        """
        if dataset_type == DatasetType.FILTER_ONLY:
            # 仅过滤
            return self.filter_and_copy_images(**kwargs)
        elif dataset_type == DatasetType.YOLO_ONLY:
            # 生成yolo格式
            return self.create_yolo_dataset(**kwargs)
        elif dataset_type == DatasetType.COCO_ONLY:
            # 生成coco格式
            return self.create_coco_dataset(**kwargs)
        elif dataset_type == DatasetType.BOTH_FORMATS:
            # 两种格式都生成
            return self.create_both_formats(**kwargs)
        else:
            raise ValueError(f"不支持的数据集类型: {dataset_type}")


def main():
    """主函数 - 支持命令行参数和直接配置两种方式"""
    parser = argparse.ArgumentParser(description="数据集处理工具")
    parser.add_argument("--type", type=str, choices=["filter_only", "yolo_only", "coco_only", "both_formats"], 
                       help="数据集类型: filter_only (仅过滤), yolo_only (生成yolo格式), coco_only (生成coco格式) 或 both_formats (两种格式都生成)")
    parser.add_argument("--images", type=str, help="图片文件夹路径")
    parser.add_argument("--labels", type=str, help="标签文件夹路径")
    parser.add_argument("--output", type=str, default="dataset_output", help="输出目录")
    parser.add_argument("--train", type=float, default=0.7, help="训练集比例 (仅yolo_only/coco_only/both_formats)")
    parser.add_argument("--val", type=float, default=0.2, help="验证集比例 (仅yolo_only/coco_only/both_formats)")
    parser.add_argument("--test", type=float, default=0.1, help="测试集比例 (仅yolo_only/coco_only/both_formats)")
    parser.add_argument("--class_names", type=str, nargs="+", help="类别名称列表（仅coco_only/both_formats需要）")
    
    args = parser.parse_args()
    
    processor = DatasetProcessor()
    
    # 如果提供了命令行参数，使用命令行模式
    if args.type:
        if args.type == "filter_only":
            if not args.images or not args.labels:
                print("错误: 仅过滤需要 --images 和 --labels 参数")
                return
            
            # 检查路径是否存在
            if not Path(args.images).exists():
                print(f"错误: 图片文件夹 {args.images} 不存在")
                return
            if not Path(args.labels).exists():
                print(f"错误: 标签文件夹 {args.labels} 不存在")
                return
            
            print("开始执行仅过滤操作...")
            processor.process_dataset(
                DatasetType.FILTER_ONLY,
                image_folder=args.images,
                label_folder=args.labels,
                output_folder=args.output
            )
             
        elif args.type == "yolo_only":
            if not args.images or not args.labels:
                print("错误: 生成yolo格式需要 --images 和 --labels 参数")
                return
            
            # 检查路径是否存在
            if not Path(args.images).exists():
                print(f"错误: 图片文件夹 {args.images} 不存在")
                return
            if not Path(args.labels).exists():
                print(f"错误: 标签文件夹 {args.labels} 不存在")
                return
            
            print("开始生成YOLO格式数据集...")
            processor.process_dataset(
                DatasetType.YOLO_ONLY,
                images_dir=args.images,
                labels_dir=args.labels,
                output_dir=args.output,
                train_ratio=args.train,
                val_ratio=args.val,
                test_ratio=args.test
            )
        
        elif args.type == "coco_only":
            if not args.images or not args.labels:
                print("错误: 生成coco格式需要 --images 和 --labels 参数")
                return
            
            # 检查路径是否存在
            if not Path(args.images).exists():
                print(f"错误: 图片文件夹 {args.images} 不存在")
                return
            if not Path(args.labels).exists():
                print(f"错误: 标签文件夹 {args.labels} 不存在")
                return
            
            print("开始生成COCO格式数据集...")
            processor.process_dataset(
                DatasetType.COCO_ONLY,
                images_dir=args.images,
                labels_dir=args.labels,
                output_dir=args.output,
                train_ratio=args.train,
                val_ratio=args.val,
                test_ratio=args.test,
                class_names=args.class_names
            )
        
        elif args.type == "both_formats":
            if not args.images or not args.labels:
                print("错误: 两种格式都生成需要 --images 和 --labels 参数")
                return
            
            # 检查路径是否存在
            if not Path(args.images).exists():
                print(f"错误: 图片文件夹 {args.images} 不存在")
                return
            if not Path(args.labels).exists():
                print(f"错误: 标签文件夹 {args.labels} 不存在")
                return
            
            print("开始同时生成YOLO和COCO格式数据集...")
            processor.process_dataset(
                DatasetType.BOTH_FORMATS,
                images_dir=args.images,
                labels_dir=args.labels,
                output_dir=args.output,
                train_ratio=args.train,
                val_ratio=args.val,
                test_ratio=args.test,
                class_names=args.class_names
            )
    else:
        # 直接配置模式 - 示例代码
        print("使用直接配置模式...")
        
        # 示例1: 仅过滤
        print("\n示例1: 仅过滤")
        processor.process_dataset(
            DatasetType.FILTER_ONLY,
            image_folder="images",
            label_folder="labelImg_data",
            output_folder="filter_only_output"
        )
        
        # 示例2: 生成yolo格式
        print("\n示例2: 生成yolo格式")
        processor.process_dataset(
            DatasetType.YOLO_ONLY,
            images_dir="images",
            labels_dir="labelImg_data",
            output_dir="yolo_only_output",
            train_ratio=0.7,
            val_ratio=0.2,
            test_ratio=0.1
        )
        
        # 示例3: 生成coco格式
        print("\n示例3: 生成coco格式")
        processor.process_dataset(
            DatasetType.COCO_ONLY,
            images_dir="images",
            labels_dir="labelImg_data",
            output_dir="coco_only_output",
            train_ratio=0.7,
            val_ratio=0.2,
            test_ratio=0.1,
            class_names=["hand_write"]
        )
        
        # 示例4: 两种格式都生成
        print("\n示例4: 两种格式都生成")
        processor.process_dataset(
            DatasetType.BOTH_FORMATS,
            images_dir="images",
            labels_dir="labelImg_data",
            output_dir="both_formats_output",
            train_ratio=0.7,
            val_ratio=0.2,
            test_ratio=0.1,
            class_names=["hand_write"]
        )


if __name__ == '__main__':
    main()